from pylab import *
from matplotlib import colors
import random

def get_color(index):           # never be white ????
    cmap = cm.jet
    norm = colors.Normalize(vmin=0,vmax=99)
    return cmap(norm(index))

def get_rand_color():
    r = random.randint(0,99)
    return get_color(r)


# def plot_bar(*item_pairs):      # all item_pairs should have the same length
# if the width or height is too small(less than 'text_not_display_ratio'), then text not display
def plot_bar(ax, group_names, item_pairs, direction="v", text_color="white", percent_color="black", text_not_display_ratio=0.03):
    assert len(group_names) == len(item_pairs[0][0])
    total = sum([pair[1] for pair in item_pairs], axis=0) # sum of each column
    max_height = max(total)
    
    ind = np.arange(len(item_pairs[0][0]))
    rects = [None] * len(item_pairs)
    bottom = np.array([0] * len(item_pairs[0][0]))
    for i,(items,values) in enumerate(item_pairs):
        if direction == "v":
            rects[i] = ax.bar(ind, values, align="center", color=get_rand_color(), bottom=bottom, width=0.5)
        else:
            rects[i] = ax.barh(ind, values, align="center", color=get_rand_color(), left=bottom, height=0.5)
        bottom += np.array(values)
        for rect,item in zip(rects[i], items):
            height = rect.get_height() if direction=="v" else rect.get_width()
            if float(height)/max_height < text_not_display_ratio: # the rect is too small, not plot text
                if float(height)/max_height < text_not_display_ratio/3:
                    continue
                if direction == "v":
                    xpos = rect.get_x() + rect.get_width()
                    ypos = rect.get_y() + rect.get_width()/2.0
                    ax.text(xpos, ypos, item, ha='left', va='bottom', weight='bold', color=text_color, rotation=30)
                else:
                    xpos = rect.get_x() + rect.get_width()/2.0
                    ypos = rect.get_y()
                    ax.text(xpos, ypos, item, ha='left', va='top', weight='bold', color=text_color, rotation=-30)
                continue
            # ypos = rect.get_y() + rect.get_height()*0.98
            xpos = rect.get_x() + rect.get_width()/2.0
            ypos = rect.get_y() + rect.get_height()/2.0
            ax.text(xpos, ypos, item, ha='center', va='center', weight='bold', color=text_color)
    # display the total percent at last
    total_height = sum(total)
    for i, rect in enumerate(rects[-1]):
        percent = total[i]*100/float(total_height)
        if direction == "v":
            ypos = rect.get_y() + rect.get_height() + total_height*0.002
            xpos = rect.get_x() + rect.get_width()/2.0
            ax.text(xpos, ypos, ("%0.1f%%" % percent), ha='center', va='bottom', weight='bold', color='black')        
        else:
            xpos = rect.get_x() + rect.get_width() + total_height*0.002
            ypos = rect.get_y() + rect.get_height()/2.0
            ax.text(xpos, ypos, ("%0.1f%%" % percent), ha='left', va='center', weight='bold', color='black')        

    

    if direction == "v":
        ax.set_xticks(ind)
        ax.set_xticklabels(group_names)
    else:
        ax.set_yticks(ind)
        ax.set_yticklabels(group_names)


        

if __name__ == "__main__":
    group_names = ["group1", "group2", "group3"]

    items1 = ( ["apple", "pear", "banana"],
               [ 10, 20, 15] )
    items2 = ( ["computer", "book", "pen"],
               # [ 7, 3, 5 ] )
               [ 500, 600, 450 ] )
    items3 = ( ["dark", "bright", "peace"],
               # [ 3, 8, 12 ] )
               [1, 2, 3] )
    # items3 = ( ["", "", "peace"],
    #            [ 0, 0, 12 ] )

    fig, ax = subplots(2,2)
    plot_bar( ax[0][0], items1[0], [items1], text_color="black")
    ax[0][0].grid(True)
    ax[0][0].set_title("bar for single array")
    
    plot_bar( ax[0][1], items1[0], [items1], "h", text_color="black")

    plot_bar( ax[1][0], group_names, [items1, items2, items3], "v", text_color="black", text_not_display_ratio=0.05)
    plot_bar( ax[1][1], group_names, [items1, items2, items3], "h", text_color="black", text_not_display_ratio=0.05)
    ax[1][1].set_title("bar for multi-arrays")
    
    show()
